from langchain_community.llms import Ollama
from langchain.callbacks.manager import CallbackManager
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
from business.utils.vector_store import VectorStore
from business.ollama.chat_doc import OllamaChat
from config.index import Config
import os
import asyncio
import json

class OllamaTest:
    def __init__(self):
        self.cfg = Config()
        self.host = self.cfg.get('ollama.host')
        self.model = self.cfg.get('ollama.model')
        
        # 初始化Ollama模型，启用流式输出
        self.llm = Ollama(
            base_url=self.host,
            model=self.model,
            callback_manager=CallbackManager([StreamingStdOutCallbackHandler()]),
            temperature=self.cfg.get('ollama.temperature'),  # 控制随机性，0-1之间，越大越随机
        )
        
        # 初始化向量存储
        self.vector_store = None
        self.init_vector_store()

    def init_vector_store(self):
        """初始化向量存储"""
        try:
            vector_store_path = os.path.abspath("data/vector")
            if not os.path.exists(vector_store_path):
                print("警告：向量库不存在，请先运行 rag_init.py 初始化向量库")
                return
                
            print("正在加载向量库...")
            vs = VectorStore()
            self.vector_store = vs.load_vector_store(vector_store_path)
            if self.vector_store:
                print("向量库加载成功！")
            else:
                print("警告：向量库加载失败")
        except Exception as e:
            print(f"加载向量库时发生错误: {e}")

    def get_relevant_context(self, query: str, k: int = 5) -> str:
        """从向量库中检索相关上下文"""
        if not self.vector_store:
            return ""
            
        try:
            # 预处理查询词
            clean_query = query.replace('-', ' ').replace('_', ' ')  # 替换特殊字符
            terms = clean_query.split()  # 分词
            
            # 构建搜索关键词
            search_terms = [clean_query]  # 原始查询（已清理）
            
            # 如果是多个词，也尝试单独搜索每个词
            if len(terms) > 1:
                search_terms.extend(terms)
            
            # 添加其他搜索变体
            for term in list(search_terms):  # 使用 list() 创建副本以便在循环中添加
                if term.strip():
                    search_terms.extend([
                        f"### {term}",  # 作为API标题搜索
                        f"{term} API",   # API相关
                        f"{term} 使用方法",  # 使用方法
                        f"{term} 示例"   # 示例代码
                    ])
            
            # 去除空字符串并去重
            search_terms = list(set(term.strip() for term in search_terms if term.strip()))
            
            all_results = []
            seen_contents = set()  # 用于去重
            
            # 对每个搜索词进行检索
            for term in search_terms:
                try:
                    results = self.vector_store.similarity_search(term, k=2)
                    for doc in results:
                        content = doc.page_content
                        # 去重并保持相关性
                        if content not in seen_contents:
                            seen_contents.add(content)
                            all_results.append(doc)
                except Exception as e:
                    print(f"搜索词 '{term}' 检索失败: {e}")
                    continue
            
            if not all_results:
                return ""
                
            # 将检索结果组合成上下文
            context_parts = []
            for i, doc in enumerate(all_results, 1):
                source = doc.metadata.get('source', '未知来源') if hasattr(doc, 'metadata') else '未知来源'
                context_parts.append(f"文档 {i} (来源: {source}):\n{doc.page_content}\n")
            
            return "\n".join(context_parts)
        except Exception as e:
            print(f"检索相关内容时发生错误: {e}")
            return ""

    def chat(self, query: str) -> str:
        """
        与模型进行对话，包含知识库检索
        :param query: 用户输入的问题
        :return: 模型的回复
        """
        try:
            # 1. 首先从向量库检索相关内容
            context = self.get_relevant_context(query)
            
            # 2. 构建提示词，包含检索到的上下文
            if context:
                print("\n找到的相关内容:")
                print("=" * 40)
                print(context)
                print("=" * 40)
                
                prompt = f"""你是一个专业的程序员助手。请基于以下知识库回答用户的问题。

参考文档:
{context}

用户问题: {query}

请使用markdown格式回答，确保代码示例使用正确的代码块标记。如果文档信息不足，请明确指出缺失的信息。"""

            else:
                prompt = f"""用户问题: {query}
                
我没有找到相关的API文档，请说明无法提供准确的API使用说明。如果你知道类似的功能，可以提供一些通用的建议。

请回答："""
            
            # 3. 调用模型生成回答
            print("\nAI助手正在思考...")
            response = self.llm.invoke(prompt)
            return response
            
        except Exception as e:
            print(f"对话过程中发生错误: {e}")
            return str(e)

async def test_stream_chat(chat: OllamaChat, query: str):
    """测试流式对话"""
    try:
        async for response in chat.stream_chat(query):
            # 直接打印流式响应的内容
            if response["event"] == "message":
                data = json.loads(response["data"])
            elif response["event"] == "error":
                print(f"\n错误: {response['data']}")
    except Exception as e:
        print(f"\n流式对话发生错误: {e}")

async def main_async():
    chat = OllamaChat()
    
    print("\n请选择测试模式:")
    print("1. 普通对话模式")
    print("2. 流式对话模式")
    mode = input("请输入模式编号 (1 或 2): ").strip()
    
    while True:
        try:
            user_input = input("\n请输入您的问题 (输入 'quit' 退出, 'switch' 切换模式): ")
            if user_input.lower() in ['quit', 'exit', 'q']:
                print("再见！")
                break
            elif user_input.lower() == 'switch':
                mode = '2' if mode == '1' else '1'
                print(f"\n已切换到{'普通' if mode == '1' else '流式'}对话模式")
                continue
                
            if user_input.strip():
                if mode == '1':
                    # 普通对话模式
                    response = chat.chat(user_input)
                    print(response)
                else:
                    # 流式对话模式
                    await test_stream_chat(chat, user_input)
                    print()  # 添加一个换行
                
        except KeyboardInterrupt:
            print("\n程序已终止")
            break
        except Exception as e:
            print(f"发生错误: {e}")

def main():
    """入口函数"""
    asyncio.run(main_async())

if __name__ == "__main__":
    main()
